module.exports = function (babel) {
  const { types: t } = babel;
  // Builder ImportDeclaration
  // Builder JSXElement


  
  return {
    name: "demo-transform", // not required
    visitor: {
      Program: {
        exit: function (path) {
          // Builder ImportDeclaration
          var importModule = t.importDeclaration([t.importDefaultSpecifier(t.identifier('ErrorBoundary'))], t.stringLiteral('@liepin/error-boundary'))
          path.unshiftContainer('body', importModule);
        }
      },
      ReturnStatement: function(path) {
        // assert type is 'JSXElement'
        if (!t.isJSXElement(path.node.argument)){
          return;
        }
        // Builder JSXElement
        var wrap = t.jSXElement(t.jSXOpeningElement(t.jSXIdentifier('ErrorBoundary'), []), t.jSXClosingElement(t.jSXIdentifier('ErrorBoundary')),[path.node.argument]);
        path.get('argument').replaceWith(wrap);
      }
    }
  };
}